import torch
import pytorch_warmup

from config import train_epoch
from .interface import Scheduler


class Cosine(Scheduler):
    def __init__(self, optimizer: torch.optim.Optimizer):
        super(Cosine, self).__init__()
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=train_epoch, eta_min=0)
        self.warmup = pytorch_warmup.LinearWarmup(optimizer, warmup_period=5)
        self.warmup.last_step = -1

    def step(self, epoch):
        self.scheduler.step(epoch)
        self.warmup.dampen()
